"""
文档处理模块
"""
import os
import sys
import re
from pathlib import Path
from typing import List, Dict, Optional
import PyPDF2
from docx import Document
import markdown

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from config import DOC_CONFIG, PDF_DIR
from utils.logger import setup_logger

logger = setup_logger(__name__)

class DocumentProcessor:
    """文档处理器，支持PDF、Word、Markdown和文本文件"""
    
    def __init__(self):
        self.supported_extensions = {
            '.pdf': self._process_pdf,
            '.txt': self._process_txt,
            '.doc': self._process_doc,
            '.docx': self._process_doc,
            '.md': self._process_markdown
        }
    
    def process_directory(self, base_dir: str = "docs") -> List[Dict]:
        """处理指定目录下的所有支持的文档"""
        if not os.path.exists(base_dir):
            os.makedirs(base_dir)
            print(f"已创建文档目录: {base_dir}")
            return []
            
        all_segments = []
        
        # 遍历目录下的所有文件
        for root, _, files in os.walk(base_dir):
            for filename in files:
                file_path = os.path.join(root, filename)
                ext = os.path.splitext(filename)[1].lower()
                
                if ext in self.supported_extensions:
                    try:
                        processor = self.supported_extensions[ext]
                        segments = processor(file_path)
                        all_segments.extend(segments)
                    except Exception as e:
                        print(f"处理文件 {filename} 时出错: {str(e)}")
                        
        return all_segments
    
    def _process_pdf(self, file_path: str) -> List[Dict]:
        """处理PDF文件"""
        segments = []
        try:
            with open(file_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)
                total_pages = len(pdf_reader.pages)
                
                for page_num in range(total_pages):
                    text = pdf_reader.pages[page_num].extract_text()
                    if text.strip():
                        segments.extend(self._split_text(
                            text,
                            os.path.basename(file_path),
                            page_num + 1
                        ))
        except Exception as e:
            print(f"PDF处理错误 {file_path}: {str(e)}")
            
        return segments
    
    def _process_txt(self, file_path: str) -> List[Dict]:
        """处理文本文件"""
        segments = []
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                text = file.read()
                if text.strip():
                    segments.extend(self._split_text(
                        text,
                        os.path.basename(file_path),
                        1  # 文本文件视为单页
                    ))
        except UnicodeDecodeError:
            # 如果UTF-8解码失败，尝试其他编码
            try:
                with open(file_path, 'r', encoding='gbk') as file:
                    text = file.read()
                    if text.strip():
                        segments.extend(self._split_text(
                            text,
                            os.path.basename(file_path),
                            1
                        ))
            except Exception as e:
                print(f"文本文件处理错误 {file_path}: {str(e)}")
                
        return segments
    
    def _process_doc(self, file_path: str) -> List[Dict]:
        """处理Word文档"""
        segments = []
        try:
            doc = Document(file_path)
            text = '\n'.join([paragraph.text for paragraph in doc.paragraphs if paragraph.text.strip()])
            if text.strip():
                segments.extend(self._split_text(
                    text,
                    os.path.basename(file_path),
                    1  # Word文档视为单页
                ))
        except Exception as e:
            print(f"Word文档处理错误 {file_path}: {str(e)}")
            
        return segments
    
    def _process_markdown(self, file_path: str) -> List[Dict]:
        """处理Markdown文件"""
        segments = []
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                md_text = file.read()
                # 转换Markdown为纯文本
                html = markdown.markdown(md_text)
                text = re.sub(r'<[^>]+>', '', html)  # 移除HTML标签
                if text.strip():
                    segments.extend(self._split_text(
                        text,
                        os.path.basename(file_path),
                        1  # Markdown文件视为单页
                    ))
        except Exception as e:
            print(f"Markdown文件处理错误 {file_path}: {str(e)}")
            
        return segments
    
    def _split_text(self, text: str, source: str, page: int) -> List[Dict]:
        """将文本分割成适当大小的片段"""
        segments = []
        paragraphs = text.split('\n\n')
        
        current_segment = ""
        for paragraph in paragraphs:
            paragraph = paragraph.strip()
            if not paragraph:
                continue
                
            # 如果当前段落加上新段落不超过最大长度，则合并
            if len(current_segment) + len(paragraph) + 1 <= 2000:
                current_segment = (current_segment + " " + paragraph).strip()
            else:
                if current_segment:
                    segments.append({
                        'text': current_segment,
                        'source': source,
                        'page': page,
                        'metadata': {'source': source, 'page': page}
                    })
                current_segment = paragraph
        
        # 添加最后一个段落
        if current_segment:
            segments.append({
                'text': current_segment,
                'source': source,
                'page': page,
                'metadata': {'source': source, 'page': page}
            })
            
        return segments
